import torch
import einops

import torch.nn as nn
from lightly.models.modules import SimCLRProjectionHead


class SignalLocalProjection(nn.Module):
    def __init__(
        self, 
        electrodes: int = 12,
        input_dim: int = 768,
        d_contrastive: int = 128,
        t_contrastive: int = 10,
        batch_size: int = 8,
        projection_interpolation: bool = True,
    ):
        super().__init__()
        
        self.local_projection = SimCLRProjectionHead(input_dim=electrodes * input_dim,
                                                     output_dim=d_contrastive)
        self.t_contrastive = t_contrastive
        self.d_contrastive = d_contrastive
        self.electrodes = electrodes
        self.batch_size = batch_size
        self.input_dim = input_dim
        self.projection_interpolation = projection_interpolation

    def forward_1(self, local_embeddings: torch.Tensor) -> torch.Tensor:
        if self.projection_interpolation:
            local_embeddings = einops.rearrange(local_embeddings,
                                                'n (c t) d -> (n t) (c d)',
                                                c=self.electrodes,
                                                n=self.batch_size)

            local_embeddings = self.local_projection(local_embeddings)
            local_embeddings = einops.rearrange(local_embeddings,
                                                '(n t) d -> n t d',
                                                d=self.d_contrastive,
                                                n=self.batch_size)
        else:
            local_embeddings = einops.rearrange(local_embeddings,
                                                'n (c t) d -> n t (c d)',
                                                c=self.electrodes,
                                                n=self.batch_size)

        return local_embeddings
    
    def forward_2(self, stacked_segments: torch.Tensor) -> torch.Tensor:
        if not self.projection_interpolation:
            stacked_segments = einops.rearrange(stacked_segments,
                                                'n t d -> (n t) d',
                                                t=self.t_contrastive,
                                                d=self.input_dim * self.electrodes) # n - number of segments for one sample
                
            stacked_segments = self.local_projection(stacked_segments)

            stacked_segments = einops.rearrange(stacked_segments,
                                                '(n t) d -> n t d',
                                                t=self.t_contrastive,
                                                d=self.d_contrastive) # n - number of segments for one sample

        return stacked_segments

class SignalGlobalProjection(nn.Module):
    def __init__(self,
                 input_dim: int = 768,
                 d_contrastive: int = 128
                 ):
        super().__init__()

        self.global_projection = SimCLRProjectionHead(input_dim=input_dim,
                                                      output_dim=d_contrastive)
        
    def forward(self, global_embeddings: torch.Tensor) -> torch.Tensor:
        global_embeddings = self.global_projection(global_embeddings)
        return global_embeddings
    